﻿using System;
using System.Collections.Generic;
using System.Data.Services.Client;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Threading;
using System.Xml;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;

namespace AntsCode.AzureStorageBackup
{ 
    public class Backup
    {
        private CloudStorageAccount account;
        private CloudStorageAccount backupToAccount;
        private string backupName;
        private int maxThreads;
        private IEnumerable<string> priorityTables;
        private IEnumerable<string> priorityContainers;

        private readonly string backupPrefix = "backup";

        public Backup(CloudStorageAccount account,
            CloudStorageAccount backupToAccount,
            int maxThreads = 5,
            IEnumerable<string> priorityTables = null,
            IEnumerable<string> priorityContainers = null)
        {
            this.account = account;
            this.backupToAccount = backupToAccount;
            this.backupName = string.Format("{0}-{1}-{2:yyyy-MM-dd-HH-mm-ss}", this.backupPrefix, account.Credentials.AccountName, DateTime.UtcNow);
            this.maxThreads = maxThreads;
            this.priorityTables = priorityTables != null ? priorityTables : new List<string>();
            this.priorityContainers = priorityContainers != null ? priorityContainers : new List<string>();
        }

        public void BeginBackup()
        {
            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Storage backup started.", DateTime.UtcNow);

            // Create the backup container
            var blobClient = StorageUtil.GetBlobClient(this.account);
            var backupToBlobClient = StorageUtil.GetBlobClient(this.backupToAccount);
            var backupToContainer = backupToBlobClient.GetContainerReference(this.backupName);
            backupToContainer.CreateIfNotExist();

            // Build the backup list
            var tables = GetTables();
            var containers = GetContainers(blobClient);
            var backupList = StorageUtil.CreateRunList(tables, containers, this.priorityTables, this.priorityContainers, BackupTable, BackupContainer);
            
            var threadRunner = new ThreadRunner();
            threadRunner.Run(backupList, this.maxThreads);

            WriteBackupData(backupToContainer, tables, containers);

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Storage backup completed.", DateTime.UtcNow);
        }

        private void WriteBackupData(CloudBlobContainer backupToContainer, IEnumerable<string> tables, IEnumerable<string> containers)
        {
            var stream = new MemoryStream();
            var writer = XmlWriter.Create(stream, new XmlWriterSettings { Indent = true});

            Trace.TraceInformation("Writing backup data xml.");

            // Add tables to manifest
            writer.WriteStartElement("BackupData");
            writer.WriteStartElement("Tables");

            foreach (var table in tables)
            {
                writer.WriteElementString("Table", table);
            }

            writer.WriteEndElement();

            // Add containers to manifest
            writer.WriteStartElement("Containers");

            foreach (var container in containers)
            {
                writer.WriteElementString("Container", container);
            }

            writer.WriteEndElement();
            writer.WriteEndElement();
            writer.Close();

            stream.SetLength(stream.Position);
            stream.Seek(0, SeekOrigin.Begin);

            var manifestBlob = backupToContainer.GetBlobReference("BackupData.xml");
            manifestBlob.UploadFromStream(stream);
        }

        private IEnumerable<string> GetContainers(CloudBlobClient blobClient)
        {            
            // Get all blob containers for the specified account
            var containers = blobClient.ListContainers().Select(c => c.Name);
           
            return containers;
        }

        private void BackupContainer(string containerName)
        {
            var blobClient = StorageUtil.GetBlobClient(this.account);
            var backupToBlobClient = StorageUtil.GetBlobClient(this.backupToAccount);

            var container = blobClient.GetContainerReference(containerName);
            var backupToContainer = backupToBlobClient.GetContainerReference(this.backupName); 
            var blobs = container.ListBlobs();
            
            string backupFileName;
            CloudBlob sourceBlob;
            CloudBlob destBlob;
            string blobName;

            if (container.Name.StartsWith(this.backupPrefix, StringComparison.CurrentCultureIgnoreCase))
            {
                // Do not backup any backup containers
                Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Container {1} excluded as it is a backup container.", DateTime.UtcNow, container.Name);
                return;
            }

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Backing up container {1} in thread {2}.", DateTime.UtcNow, container.Name, Thread.CurrentThread.ManagedThreadId);

            foreach (var blob in blobs)
            {
                // Get the full blob name including container
                blobName = blob.Uri.AbsolutePath.Substring(blob.Uri.AbsolutePath.IndexOf(container.Name));
               
                try
                {
                    backupFileName = string.Format("BlobData/{0}", blobName);
                    sourceBlob = container.GetBlobReference(blob.Uri.ToString());
                    destBlob = backupToContainer.GetBlobReference(backupFileName);

                    StorageUtil.CopyBlob(sourceBlob, destBlob);
                }
                catch (Exception e)
                {
                    // The blob may have been deleted since retrieving the blob list so ignore
                    Trace.TraceWarning("{0:yyyy MM dd hh:mm:ss} Backup of blob {1} has failed with error {2}", DateTime.UtcNow, blobName, e.Message);
                }
            }
                
            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Finished backing up container {1}", DateTime.UtcNow, container.Name);
        }
       
        private IEnumerable<string> GetTables()
        {
            // Get all tables for the specified account
            var tableClient = StorageUtil.GetTableClient(this.account);
            var tableNames = tableClient.ListTables().ToList();
            
            return tableNames;
        }
 
        private void BackupTable(string tableName)
        {
            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Starting to backup table {1} in thread {2}", DateTime.UtcNow, tableName, Thread.CurrentThread.ManagedThreadId);

            var tableClient = StorageUtil.GetTableClient(this.account);
            var context = tableClient.GetDataServiceContext();
            
            context.MergeOption = MergeOption.NoTracking;
            context.ResolveType = Backup.ResolveType;
            context.ReadingEntity += new EventHandler<ReadingWritingEntityEventArgs>(Backup.OnReadingEntity);
            context.RetryPolicy = RetryPolicies.RetryExponential(5, RetryPolicies.DefaultClientBackoff);

            var query = from entity in context.CreateQuery<BackupEntity>(tableName) select entity;
            var cloudQuery = new CloudTableQuery<BackupEntity>((DataServiceQuery<BackupEntity>)query);

            WriteToBlob(tableName, cloudQuery);

            Trace.TraceInformation("{0:yyyy MM dd hh:mm:ss} Finished backup of table {1}", DateTime.UtcNow, tableName);
        }

        private void WriteToBlob(string tableName, CloudTableQuery<BackupEntity> query)
        {
            // A block can be at most 4 MB in Azure Storage. Though we will be using much less
            // we will allocate 4MB for the edge case where an entity may be 1MB
            MemoryStream stream = new MemoryStream(4 * 1024 * 1024);
            State state = new State(query, stream);
            CloudBlockBlob tempBlob;
            CloudBlockBlob backupBlob;
            string backupFileName;
            var backupToBlobClient = StorageUtil.GetBlobClient(this.backupToAccount);
            var backupToContainer = backupToBlobClient.GetContainerReference(this.backupName);
            TableServiceEntity firstEntity = null;
            TableServiceEntity lastEntity = null;

            while (!state.HasCompleted)
            {
                backupFileName = string.Format("TableData/{0}/{1}.xml",
                    tableName,
                    Guid.NewGuid());

                tempBlob = backupToContainer.GetBlockBlobReference(backupFileName);

                Blob blob = new Blob(state);

                List<string> blockIdList = new List<string>();

                foreach (Block block in blob.Blocks)
                {
                    string blockId = this.BackupBlock(stream, tempBlob, block, out firstEntity, out lastEntity);
                    if (!string.IsNullOrEmpty(blockId))
                    {
                        blockIdList.Add(blockId);
                    }
                }

                if (blockIdList.Count > 0)
                {
                    // Commit block list
                    tempBlob.PutBlockList(blockIdList);

                    if (firstEntity != null && lastEntity != null)
                    {
                        backupFileName = string.Format("TableData/{0}/{1}-{2}-{3}-{4}",
                            tableName,
                            firstEntity.PartitionKey,
                            firstEntity.RowKey,
                            lastEntity.PartitionKey,
                            lastEntity.RowKey);

                        if (backupFileName.Length > 252)
                        {
                            backupFileName = backupFileName.Substring(0, 252);
                        }

                        backupBlob = backupToContainer.GetBlockBlobReference(backupFileName + ".xml");
                        backupBlob.CopyFromBlob(tempBlob);
                    }

                    tempBlob.Delete();
                }
            }
        }

        private string BackupBlock(MemoryStream stream, CloudBlockBlob backupBlob, Block block, out TableServiceEntity firstEntity, out TableServiceEntity lastEntity)
        {
            int entityCount = 0;
            firstEntity = null;
            lastEntity = null;

            // reset the memory stream as we begin a new block
            stream.Seek(0, SeekOrigin.Begin);
            stream.SetLength(0);
            XmlWriter writer = XmlWriter.Create(stream);

            writer.WriteStartElement("Block");
            foreach (Batch batch in block.Batches)
            {
                // write begin batch statement
                writer.WriteStartElement("Batch");
                foreach (BackupEntity entity in batch.Entities)
                {
                    entityCount++;
                    entity.EntryElement.WriteTo(writer);

                    if (firstEntity == null)
                    {
                        firstEntity = entity;
                    }

                    lastEntity = entity;
                }
                writer.WriteEndElement();

            }
            writer.WriteEndElement();
            writer.Close();
            stream.SetLength(stream.Position);
            stream.Seek(0, SeekOrigin.Begin);

            // if we have written > 0 entities, let us store to a block. Else we can reject this block
            if (entityCount > 0)
            {
                backupBlob.PutBlock(block.BlockId, stream, null);
                return block.BlockId;
            }

            return null;
        }

        /// <summary>
        /// The class that maintains the global state for the iteration
        /// </summary>
        internal class State
        {
            protected MemoryStream stream;
            IEnumerator<BackupEntity> queryIterator;

            internal State(CloudTableQuery<BackupEntity> query, MemoryStream stream)
            {
                this.queryIterator = query.GetEnumerator();
                this.stream = stream;
            }

            /// <summary>
            /// This entity is the one we may have retrieved but it does not belong to the batch
            /// So we store it here so that it can be returned on the next iteration
            /// </summary>
            internal BackupEntity LookAheadEntity { private get; set; }

            /// <summary>
            /// We have completed if look ahead entity is null and iterator is completed too.
            /// </summary>
            internal bool HasCompleted
            {
                get
                {
                    return this.queryIterator == null && this.LookAheadEntity == null;
                }
            }

            /// <summary>
            /// Get the amount of data we have saved in the entity
            /// </summary>
            internal long CurrentBlockSize
            {
                get
                {
                    stream.Flush();
                    return stream.Position;
                }
            }

            /// <summary>
            /// Return the next entity - which can be either the 
            /// look ahead entity or a new one from the iterator.
            /// We return null if there are no more entities
            /// </summary>
            /// <returns></returns>
            internal BackupEntity GetNextEntity()
            {
                BackupEntity entityToReturn = null;
                if (this.LookAheadEntity != null)
                {
                    entityToReturn = this.LookAheadEntity;
                    this.LookAheadEntity = null;
                }
                else if (this.queryIterator != null)
                {
                    if (this.queryIterator.MoveNext())
                    {
                        entityToReturn = this.queryIterator.Current;
                    }
                    else
                    {
                        this.queryIterator = null;
                    }
                }

                return entityToReturn;
            }
        }

        /// <summary>
        /// Represents a collection of entities in a single batch
        /// </summary>
        internal class Batch
        {
            static int MaxEntityCount = 100;
            // Save at most 3.5MB in a batch so that we have enough room for 
            // the xml tags that WCF Data Services adds in the OData protocol
            static int MaxBatchSize = (int)(3.5 * 1024 * 1024);

            State state;

            internal Batch(State state)
            {
                this.state = state;
            }

            /// <summary>
            /// Yield entities until we hit a condition that should terminate a batch.
            /// The conditions to terminate on are:
            /// 1. 100 entities in a batch
            /// 2. 3.5MB of data
            /// 2. 3.8MB of block size
            /// 3. We see a new partition key
            /// </summary>
            internal IEnumerable<BackupEntity> Entities
            {
                get
                {
                    BackupEntity entity;
                    long currentSize = this.state.CurrentBlockSize;

                    string lastPartitionKeySeen = null;
                    int entityCount = 0;

                    while ((entity = state.GetNextEntity()) != null)
                    {
                        if (lastPartitionKeySeen == null)
                        {
                            lastPartitionKeySeen = entity.PartitionKey;
                        }

                        int approxEntitySize = entity.EntryElement.ToString().Length * 2;
                        long batchSize = this.state.CurrentBlockSize - currentSize;
                        if (entityCount >= Batch.MaxEntityCount
                            || !string.Equals(entity.PartitionKey, lastPartitionKeySeen)
                            || batchSize + approxEntitySize > Batch.MaxBatchSize
                            || this.state.CurrentBlockSize + approxEntitySize > Block.MaxBlockSize)
                        {
                            // set this current entity as the look ahead since it needs to be part of the next batch
                            state.LookAheadEntity = entity;                      
                            yield break;
                        }

                        entityCount++;
                        yield return entity;
                    }
                }
            }
        }

        /// <summary>
        /// Represents all batches in a block
        /// </summary>
        internal class Block
        {
            // Though a block can be of 4MB we will stop before to allow buffer
            public static int MaxBlockSize = (int)(3.8 * 1024 * 1024);

            State state;

            internal string BlockId { get; private set; }

            internal Block(State state)
            {
                this.state = state;
                this.BlockId = Convert.ToBase64String(Guid.NewGuid().ToByteArray());
            }

            /// <summary>
            /// The list of batches in the block. 
            /// </summary>
            internal IEnumerable<Batch> Batches
            {
                get
                {
                    while (!state.HasCompleted && state.CurrentBlockSize < Block.MaxBlockSize)
                    {
                        yield return new Batch(state);
                    }
                }
            }
        }

        /// <summary>
        /// Represents all blocks in a blob
        /// </summary>
        internal class Blob
        {
            /// <summary>
            /// We will allow storing at most 20 blocks in a blob
            /// </summary>
            static int MaxBlocksInBlobs = 20;

            State state;
            internal CloudBlob blob { get; private set; }

            internal Blob(State state)
            {
                this.state = state;
            }

            /// <summary>
            /// The blocks that form the blob
            /// </summary>
            internal IEnumerable<Block> Blocks
            {
                get
                {
                    int blockCount = 0;

                    while (!state.HasCompleted && blockCount < Blob.MaxBlocksInBlobs)
                    {
                        blockCount++;
                        yield return new Block(state);
                    }
                }
            }
        }

        private static Type ResolveType(string entityName)
        {
            return typeof(BackupEntity);
        }

        private static void OnReadingEntity(object sender, ReadingWritingEntityEventArgs args)
        {
            BackupEntity entity = args.Entity as BackupEntity;
            entity.EntryElement = args.Data;
        }
    }
}